import json
import os
from pathlib import Path

def process_json_file(input_file, output_dir):
    """
    处理JSON文件，提取指定字段并按dataset分类保存
    :param input_file: 输入JSON文件路径
    :param output_dir: 输出目录路径
    """
    try:
        with open(input_file, 'r', encoding='utf-8') as f:
            data = json.load(f)
        
        # 创建三个子目录
        datasets = ['train', 'val', 'test']
        for dataset in datasets:
            os.makedirs(os.path.join(output_dir, dataset), exist_ok=True)
        
        # 初始化统计数据
        stats = {dataset: 0 for dataset in datasets}
        prompt = "请根据提供的标签内容对图片进行分类："
        
        for item in data:
            dataset = item.get("dataset", "train")  # 默认为train
            if dataset not in datasets:
                dataset = "train"  # 如果不是train/val/test，默认归为train
                
            # 处理tag字段：替换|为,
            processed_tag = item.get("tag", "").replace("|", ",")
            
            # 按指定顺序构建新对象
            processed_item = {
                "image": item.get("image", ""),
                "prompt": prompt,
                "tag": processed_tag,
                "父级分类名称": item.get("父级分类名称", ""),
                "分类名称": item.get("分类名称", "")
            }
            
            # 保存到对应的dataset文件
            output_file = os.path.join(output_dir, dataset, f"data_classify_{dataset}.json")
            
            # 读取已有数据或创建新文件
            if os.path.exists(output_file):
                with open(output_file, 'r', encoding='utf-8') as f:
                    existing_data = json.load(f)
            else:
                existing_data = []
            
            existing_data.append(processed_item)
            
            # 写入文件（保持字段顺序）
            with open(output_file, 'w', encoding='utf-8') as f:
                json.dump(existing_data, f, ensure_ascii=False, indent=2, sort_keys=False)
            
            stats[dataset] += 1
        
        print(f"\n处理完成！数据分布统计：")
        for dataset, count in stats.items():
            print(f"{dataset}集: {count}条 → {output_dir}/{dataset}/data_classify_{dataset}.json")
        print("字段顺序保持：image → prompt → tag → 父级分类名称 → 分类名称")
    
    except FileNotFoundError:
        print(f"错误：输入文件 {input_file} 不存在")
    except json.JSONDecodeError:
        print("错误：输入文件不是有效的JSON格式")
    except Exception as e:
        print(f"处理过程中发生错误：{str(e)}")

# 使用示例
if __name__ == "__main__":
    input_json = "final_dataset.json"  # 替换为您的输入文件路径
    output_dir = "processed_data_classify"     # 输出目录路径
    process_json_file(input_json, output_dir)